#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Author : gao
# Time : 2020/7/6 14:34

import Shannon as sn
import numpy as np

'''
    本模块为决策树选取数据
'''

'''
    功能 : 切分数据集
    参数 ：
        dataSet : 要划分的数据集
        axis : 数据的第n维属性
        value ：划分第axis维的数据的目标值
    返回值：第axis维等于value的矩阵
'''


def splitDataSet(dataSet: list, axis, value) -> list:
    valueList = []
    for data in dataSet:
        if data[axis] == value:
            dataNew = data[:axis]
            dataNew.extend(data[axis + 1:])
            valueList.append(dataNew)
    return valueList


'''
    根据信息增益
            (InfoGain =ShannonEntropy - 求和 该属性的各个值对应的比例*该值对应的数据集的香浓熵 )
            来选取最好的属性对应的下标
'''


def selectAxis(dataSet: list):
    # 数据集的香农熵
    entDataSet = sn.calcShannonEnt(dataSet)
    # 属性数，最后一个是标签
    features = len(dataSet[0]) - 1
    # 信息增益
    infoGain = 0.0
    # 最优的属性 初始为-1
    bestFeature = -1
    # 开始遍历，找最大的信息增益
    for i in range(features):
        # 获取属性值集合
        featureValueSet = set([d[i] for d in dataSet])
        # 总数居量
        numTotalData = len(dataSet)
        # 暂存的信息增益
        tmpGain = entDataSet
        # 遍历切割数据集，然后计算信息熵
        for f in featureValueSet:
            # 切割数据集
            dataSetBySplit = splitDataSet(dataSet, i, f)
            # 数据集的长度
            numSplitData = len(dataSetBySplit)
            # 计算信息增益
            tmpGain -= numSplitData / numTotalData * sn.calcShannonEnt(dataSetBySplit)
        print('第%d个属性的信息增益为%.3f' % (i, tmpGain))
        if tmpGain > infoGain:
            bestFeature = i
            infoGain = tmpGain
    return bestFeature


if __name__ == '__main__':
    # dataSet, labels = sn.createDataSet()
    # print('最优的属性的索引是', str(selectAxis(dataSet)))
    results=splitDataSet([[0,1],[1,3],[1,3]],1,3)
    print(results)
